package kvraft

import (
	"6.824/labgob"
	"6.824/labrpc"
	"6.824/raft"
	"bytes"
	"fmt"
	"log"
	"math"
	"sync"
	"sync/atomic"
	"time"
)

const Debug = false

func DPrintf(format string, a ...interface{}) (n int, err error) {
	if Debug {
		log.Printf(format, a...)
	}
	return
}

type Op struct {
	// Your definitions here.
	// Field names must start with capital letters,
	// otherwise RPC will break.
	Op        string
	Key       string
	Value     string
	ClientId  int64
	CommandId int
}

const TIMEOUT = 500 * time.Millisecond

type CommandResult CommonReply

type KVServer struct {
	mu           sync.Mutex
	me           int
	rf           *raft.Raft
	applyCh      chan raft.ApplyMsg
	dead         int32 // set by Kill()
	maxraftstate int   // snapshot if log grows this big

	// Your definitions here.

	database           map[string]string
	lastIncludeIndex   int
	waitChMap          map[int]chan CommandResult // server 提交日志时需要等待日志在 Raft 层同步完毕，key 为日志索引，val 为 server 陷入等待的 ch，这等效于 cond
	executedCommandMap map[int64]int              // key 是客户端 ID，val 是客户端成功执行的最大的 commandId(非 Get)，单独的客户端总是串行执行命令的
}

// getWaitCh 获取 waitChMap  中 key 为 index 的 chan，如果没有则创建，方法 unsafe
func (kv *KVServer) getWaitCh(index int) chan CommandResult {
	ch, contains := kv.waitChMap[index]
	if !contains {
		ch = make(chan CommandResult, 1)
		kv.waitChMap[index] = ch
	}
	return ch
}

// apply 应用到数据库中 ，方法 unsafe
func (kv *KVServer) apply(op Op) *CommandResult {
	result := &CommandResult{}
	defer DPrintf("%s apply op %+v and result is %+v", kv.rf.Self(), op, result)

	result.Err = OK

	switch op.Op {
	case "Get":
		val, contains := kv.database[op.Key]
		if contains {
			result.Value = val
		} else {
			result.Err = ErrNoKey
		}
	case "Put":
		kv.database[op.Key] = op.Value
		result.Value = kv.database[op.Key]
	case "Append":
		kv.database[op.Key] += op.Value
		result.Value = kv.database[op.Key]
	}
	return result
}

func (kv *KVServer) applyTicker() {
	for true {
		select {
		case applyMsg := <-kv.applyCh:
			if applyMsg.CommandValid {
				kv.mu.Lock()
				if applyMsg.CommandIndex <= kv.lastIncludeIndex {
					kv.mu.Unlock()
					continue
				}
				kv.lastIncludeIndex = applyMsg.CommandIndex

				op := applyMsg.Command.(Op)
				result := &CommandResult{}

				DPrintf("%s apply message %+v, commandId is %d, maxCommandId is %d\n", kv.rf.Self(), applyMsg, op.CommandId, kv.executedCommandMap[op.ClientId])

				// Get 请求幂等不去重，其他请求去重，只有 op.CommandId > kv.executedCommandMap[op.ClientId] 才执行
				if op.Op == "Get" || op.CommandId > kv.executedCommandMap[op.ClientId] {
					result = kv.apply(op)
					kv.executedCommandMap[op.ClientId] = int(math.Max(float64(op.CommandId), float64(kv.executedCommandMap[op.ClientId])))
				}

				DPrintf("%s now log on key %s is %s\n", kv.rf.Self(), op.Key, kv.database[op.Key])

				if kv.maxraftstate != -1 && kv.rf.RaftStateSize() >= kv.maxraftstate {
					DPrintf("Start Snapshot\n\n")
					kv.rf.Snapshot(applyMsg.CommandIndex, kv.snapshotData())
				}

				// 不是 Leader 则不通知，避免创建大量无效对象
				if _, isLeader := kv.rf.GetState(); isLeader {
					kv.getWaitCh(applyMsg.CommandIndex) <- *result
				}
				kv.mu.Unlock()
			} else if applyMsg.SnapshotValid {
				// 这是一条快照信息，更新 lastIncludeIndex，读取快照数据恢复
				kv.mu.Lock()
				kv.readSnapshot(applyMsg.Snapshot)
				kv.lastIncludeIndex = int(math.Max(float64(kv.lastIncludeIndex), float64(applyMsg.SnapshotIndex)))
				kv.mu.Unlock()
			}
		case <-time.After(1 * time.Second):
			if kv.killed() {
				return
			}
		}
	}
}

func (kv *KVServer) Get(args *GetArgs, reply *CommonReply) {
	// Your code here.
	if _, isLeader := kv.rf.GetState(); !isLeader {
		reply.Err = ErrWrongLeader
		return
	}

	op := Op{Op: "Get", CommandId: args.CommandId, Key: args.Key}
	index, _, isLeader := kv.rf.Start(op)

	if !isLeader {
		reply.Err = ErrWrongLeader
		return
	}

	kv.mu.Lock()
	ch := kv.getWaitCh(index)
	kv.mu.Unlock()

	defer func() {
		kv.mu.Lock()
		delete(kv.waitChMap, index)
		kv.mu.Unlock()
	}()

	select {
	case result := <-ch:
		reply.Value, reply.Err = result.Value, result.Err
	case <-time.After(TIMEOUT):
		reply.Err = ErrTimeOut
	}
}

func (kv *KVServer) PutAppend(args *PutAppendArgs, reply *CommonReply) {
	// Your code here.
	if _, isLeader := kv.rf.GetState(); !isLeader {
		reply.Err = ErrWrongLeader
		return
	}

	op := Op{Op: args.Op, CommandId: args.CommandId, Key: args.Key, Value: args.Value, ClientId: args.ClientId}

	kv.mu.Lock()
	if cid, contains := kv.executedCommandMap[op.ClientId]; contains && cid >= op.CommandId {
		reply.Err = OK
		kv.mu.Unlock()
		return
	}
	kv.mu.Unlock()

	index, _, _ := kv.rf.Start(op)

	kv.mu.Lock()
	ch := kv.getWaitCh(index)
	kv.mu.Unlock()

	defer func() {
		kv.mu.Lock()
		delete(kv.waitChMap, index)
		kv.mu.Unlock()
	}()

	select {
	case result := <-ch:
		reply.Err, reply.Value = result.Err, result.Value
	case <-time.After(TIMEOUT):
		reply.Err = ErrTimeOut
	}
}

//
// the tester calls Kill() when a KVServer instance won't
// be needed again. for your convenience, we supply
// code to set rf.dead (without needing a lock),
// and a killed() method to test rf.dead in
// long-running loops. you can also add your own
// code to Kill(). you're not required to do anything
// about this, but it may be convenient (for example)
// to suppress debug output from a Kill()ed instance.
//
func (kv *KVServer) Kill() {
	atomic.StoreInt32(&kv.dead, 1)
	kv.rf.Kill()
	// Your code here, if desired.
}

func (kv *KVServer) killed() bool {
	z := atomic.LoadInt32(&kv.dead)
	return z == 1
}

func (kv *KVServer) snapshotData() []byte {
	w := new(bytes.Buffer)
	e := labgob.NewEncoder(w)
	e.Encode(kv.database)
	e.Encode(kv.executedCommandMap)
	data := w.Bytes()
	return data
}

func (kv *KVServer) readSnapshot(snapshot []byte) {
	if snapshot == nil || len(snapshot) == 0 {
		return
	}

	r := bytes.NewBuffer(snapshot)
	d := labgob.NewDecoder(r)

	var database map[string]string
	var executedCommandMap map[int64]int

	if d.Decode(&database) == nil && d.Decode(&executedCommandMap) == nil {
		kv.database = database
		kv.executedCommandMap = executedCommandMap
	} else {
		fmt.Printf("Read Error")
	}
}

//
// servers[] contains the ports of the set of
// servers that will cooperate via Raft to
// form the fault-tolerant key/value service.
// me is the index of the current server in servers[].
// the k/v server should store snapshots through the underlying Raft
// implementation, which should call persister.SaveStateAndSnapshot() to
// atomically save the Raft state along with the snapshot.
// the k/v server should snapshot when Raft's saved state exceeds maxraftstate bytes,
// in order to allow Raft to garbage-collect its log. if maxraftstate is -1,
// you don't need to snapshot.
// StartKVServer() must return quickly, so it should start goroutines
// for any long-running work.
//
func StartKVServer(servers []*labrpc.ClientEnd, me int, persister *raft.Persister, maxraftstate int) *KVServer {
	// call labgob.Register on structures you want
	// Go's RPC library to marshall/unmarshall.
	labgob.Register(Op{})

	kv := new(KVServer)
	kv.me = me
	kv.maxraftstate = maxraftstate

	kv.applyCh = make(chan raft.ApplyMsg)
	kv.rf = raft.Make(servers, me, persister, kv.applyCh)

	// You may need initialization code here.

	kv.executedCommandMap = make(map[int64]int)
	kv.database = make(map[string]string)
	kv.waitChMap = make(map[int]chan CommandResult)
	kv.lastIncludeIndex = 0
	kv.readSnapshot(persister.ReadSnapshot())

	go kv.applyTicker()
	return kv
}
